package org.sunyaxing.transflow.agent.agentplugin;

import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.matcher.ElementMatchers;

import java.io.PrintStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

public class TransflowAgent {
    private static PrintStream ps = System.out;

    public static void agentmain(String agentArgs, Instrumentation instrumentation) {
        Map<String, String> args = parse(agentArgs);
        String model = args.getOrDefault("model", "install");
        if ("uninstall".equals(model)) {
            if (TransflowBootstrap.installed) {
                TransflowBootstrap.destory();
            }
        } else {
            String packages = args.getOrDefault("packages", "");
            if (packages == null || packages.isEmpty()) {
                return;
            }
            if (TransflowBootstrap.installed) {
                return;
            }
            TransflowBootstrap.getInstance(instrumentation);
            ClassFileTransformer classFileTransformer = new AgentBuilder.Default()
                    .with(AgentBuilder.InitializationStrategy.NoOp.INSTANCE)
                    .with(AgentBuilder.RedefinitionStrategy.REDEFINITION)
                    .type(
                            ElementMatchers.nameStartsWith(packages)
                                    .and(ElementMatchers.not(ElementMatchers.nameStartsWith("org.sunyaxing.transflow")))
                                    .and(ElementMatchers.not(ElementMatchers.nameContains("$")))
                                    .and(ElementMatchers.not(ElementMatchers.isInterface()))
                    )
                    .transform((builder, typeDescription, classLoader, module) -> {
                        // 备份原始字节码
                        OrgClassByteManager.put(typeDescription.getName(), classLoader, builder.make().getBytes());
                        DynamicType.Builder<?> dynamicType = builder
                                .visit(Advice.to(ControllerAdvice.class).on(
                                        ElementMatchers.isMethod()
                                                .and(ElementMatchers.not(ElementMatchers.isConstructor()))
                                                .and(ElementMatchers.not(ElementMatchers.isStatic()))
                                                .and(ElementMatchers.not(ElementMatchers.nameStartsWith("get")))
                                                .and(ElementMatchers.not(ElementMatchers.nameStartsWith("set")))
                                ));
                        return dynamicType;
                    }).installOn(instrumentation);
            TransflowBootstrap.setClassFileTransformer(classFileTransformer);
        }
    }

    public static Map<String, String> parse(String agentArgs) {
        Map<String, String> params = new HashMap<>();
        Optional.ofNullable(agentArgs)
                .ifPresent(args -> {
                    for (String pair : args.split(",")) {
                        String[] kv = pair.split("=", 2);
                        if (kv.length == 2) {
                            params.put(kv[0].trim(), kv[1].trim());
                        }
                    }
                });
        return params;
    }
}
